#include <vector>
#include<iostream>
#include<stdio.h>
#include <algorithm>
#include "softmax_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
using namespace std;


namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    TilingData tiling;
    uint32_t sizeofdatatype;
    uint32_t totalLengthAligned;

    // 1. 获取平台信息
    uint64_t ub_size;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);

    // 2. 获取数据信息
    int64_t dim = *(context->GetAttrs()->GetInt(0));
    
    auto shape = context->GetInputTensor(0)->GetOriginShape();
    const int32_t dimNum = shape.GetDimNum();

    uint32_t typeLength = 0;
    ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);

    if(dim<0) dim = dimNum+dim;

    // if (dim!=dimNum-1) exit(-1);
    int32_t batch_size = 1;
    int32_t batch_number = 1;
    int32_t stride_size = 1;
    int32_t stride_number = 1;
    int32_t reduce_size = 1;
    int32_t reduce_number = 1;
    int32_t total_size = 1;

    for (int32_t i = 0; i<dimNum; ++i)
    {
        int32_t shape_dim = shape.GetDim(i);
        if (i<dim)
        {
            batch_number*=shape_dim;
        }
        else if (i==dim)
        {
            reduce_number = shape_dim;
            batch_size*=shape_dim;
        }
        else
        {
            batch_size*=shape_dim;
            stride_number*=shape_dim;
            reduce_size*=shape_dim;
        }

        total_size*=shape_dim;
    }

    tiling.set_total_size(total_size);
    tiling.set_batch_size(batch_size);
    tiling.set_batch_number(batch_number);
    tiling.set_stride_size(stride_size);
    tiling.set_stride_number(stride_number);
    tiling.set_reduce_size(reduce_size);
    tiling.set_reduce_number(reduce_number);

    if ((reduce_number*typeLength)%32!=0) tiling.set_mode(0);
    else {
        tiling.set_mode(1);
        if (reduce_number>20000) exit(-1);
    }


    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                        context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t* currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
static graphStatus InferDataType(gert::InferDataTypeContext* context)
{
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
}

namespace ops {
class Softmax : public OpDef {
public:
    explicit Softmax(const char* name) : OpDef(name)
    {
        this->Input("x")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        // this->Input("axes")
        //         .ParamType(REQUIRED)
        //         .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
        //         .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
        //         .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->Output("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("dim").AttrType(OPTIONAL).Int(int64_t(-1));
        
        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend310b");
    }
};
OP_ADD(Softmax);
}
